import numpy
import torch
import matplotlib.pyplot as plot


def multi_box(height, width, sizes, ratios, device):  # height,width: 高宽步长
    steps_h = 1.0 / height
    steps_w = 1.0 / width
    n = len(sizes) + len(ratios) - 1  # n = len(sizes) + len(ratios) - 1 : 每个像素要生成的锚框的数量
    center_h = (torch.arange(height, device=device) + 0.5) * steps_h
    center_w = (torch.arange(width, device=device) + 0.5) * steps_w

    y, x = torch.meshgrid(center_h, center_w)
    y = y.reshape(-1)
    x = x.reshape(-1)
    s = torch.tensor(sizes, device=device)
    r = torch.tensor(ratios, device=device)
    w = torch.cat([(s * torch.sqrt(r[0])), s[0] * torch.sqrt(r[1:])]) * height / width
    h = torch.cat([(s / torch.sqrt(r[0])), s[0] / torch.sqrt(r[1:])])
    offsets = torch.stack([-w, -h, w, h], dim=1).repeat(width * height, 1) / 2.0
    centers = torch.stack([x, y, x, y], dim=1).repeat_interleave(n, dim=0)
    return offsets + centers


# 生成的锚框：[x1,y1,x2,y2]

def bbox_rectangle(self, bbox, color):
    return plot.Rectangle(
        xy=(bbox[0], bbox[1]), width=bbox[2] - bbox[0], height=bbox[3] - bbox[1],
        fill=False, edgecolor=color, linewidth=2
    )


def show_boxes(figure, boxes, width, height, color):
    boxes = boxes * torch.tensor([width, height, width, height])

    for i, bbox in enumerate(boxes):
        bbox = bbox.detach().numpy()
        rect = plot.Rectangle(
            xy=(bbox[0], bbox[1]), width=bbox[2] - bbox[0], height=bbox[3] - bbox[1],
            fill=False, edgecolor=color, linewidth=2
        )
        figure.axes.add_patch(rect)


@torch.no_grad()
def iou(boxes_o_1, boxes_o_2, mode='joint'):  # 计算两个框之间的交并比 input shape: [boxes_num,4]

    boxes_1 = boxes_o_1.reshape(-1, boxes_o_1.shape[-1])
    boxes_2 = boxes_o_2.reshape(-1, boxes_o_2.shape[-1])

    areas_1 = (boxes_1[:, 2] - boxes_1[:, 0]) * (boxes_1[:, 3] - boxes_1[:, 1])
    areas_2 = (boxes_2[:, 2] - boxes_2[:, 0]) * (boxes_2[:, 3] - boxes_2[:, 1])
    if mode == 'joint':
        inter_up_left = torch.max(boxes_1[:, None, 0:2], boxes_2[:, 0:2])
        inter_down_right = torch.min(boxes_1[:, None, 2:], boxes_2[:, 2:])
        inters = (inter_down_right - inter_up_left).clamp(min=0)  # 将小于0的值都设为0
        inter_areas = inters[:, :, 0] * inters[:, :, 1]
        union_areas = areas_1[:, None] + areas_2 - inter_areas
        r = inter_areas / union_areas  # table shape: boxes_1.shape[0] * boxes_2.shape[0]
    else:
        inter_up_left = torch.max(boxes_1[:, 0:2], boxes_2[:, 0:2])
        inter_down_right = torch.min(boxes_1[:, 2:], boxes_2[:, 2:])
        inters = (inter_down_right - inter_up_left).clamp(min=0)  # 将小于0的值都设为0
        inter_areas = inters[:, 0] * inters[:, 1]
        union_areas = areas_1 + areas_2 - inter_areas
        r = inter_areas / union_areas
    # return r.view(boxes_1.shape[0], boxes_2.shape[0])
    return r


# anchor_bboxes,true_bboxes shape: [:,4]
def anchor_true_offset(anchor_bboxes, true_bboxes, offsets=None):  # 计算锚框和真实框之间的偏移量；根据锚框和偏移量来获得真实框
    mu_x, mu_y, mu_w, mu_h, sigma_x, sigma_y, sigma_w, sigma_h = 0, 0, 0, 0, 0.1, 0.1, 0.2, 0.2
    eps = 1e-6  # 防止下面的log函数参数为0的情况
    if offsets is None:
        anchor_bboxes_c = bbox_position_cast(anchor_bboxes, 1)
        true_bboxes_c = bbox_position_cast(true_bboxes, 1)
        offset_x = ((true_bboxes_c[:, 0] - anchor_bboxes_c[:, 0]) / anchor_bboxes_c[:, 2] - mu_x) / sigma_x
        offset_y = ((true_bboxes_c[:, 1] - anchor_bboxes_c[:, 1]) / anchor_bboxes_c[:, 3] - mu_y) / sigma_y
        offset_w = (torch.log(true_bboxes_c[:, 2].clamp(min=eps) / anchor_bboxes_c[:, 2]) - mu_w) / sigma_w
        offset_h = (torch.log(true_bboxes_c[:, 3].clamp(min=eps) / anchor_bboxes_c[:, 3]) - mu_h) / sigma_h
        return torch.stack([offset_x, offset_y, offset_w, offset_h], dim=1)
        # offset_x,offset_y:中心点偏差,offset_w,offset_h:宽，高偏差
    else:
        anchor_boxes_c = bbox_position_cast(anchor_bboxes, 1)
        true_boxes_x = ((offsets[:, 0] * sigma_x) + mu_x) * anchor_boxes_c[:, 2] + anchor_boxes_c[:, 0]
        true_boxes_y = ((offsets[:, 1] * sigma_y) + mu_y) * anchor_boxes_c[:, 3] + anchor_boxes_c[:, 1]
        true_boxes_w = torch.exp((offsets[:, 2] * sigma_w) + mu_w) * anchor_boxes_c[:, 2]
        true_boxes_h = torch.exp((offsets[:, 3] * sigma_h) + mu_h) * anchor_boxes_c[:, 3]
        true_boxes = torch.stack([true_boxes_x, true_boxes_y, true_boxes_w, true_boxes_h], dim=1)
        return bbox_position_cast(true_boxes, 0)


# 不改变原box的任何信息
def bbox_position_cast(boxes, which):
    boxes_0 = boxes.reshape(-1, boxes.shape[-1])

    if which:
        x1, y1, x2, y2 = boxes_0[:, 0], boxes_0[:, 1], boxes_0[:, 2], boxes_0[:, 3]
        center_x = (x1 + x2) / 2
        center_y = (y1 + y2) / 2
        width = x2 - x1
        height = y2 - y1
        boxes_c = torch.stack([center_x, center_y, width, height], dim=1)
    else:
        center_x, center_y, width, height = boxes_0[:, 0], boxes_0[:, 1], boxes_0[:, 2], boxes_0[:, 3]
        x1 = center_x - 0.5 * width
        y1 = center_y - 0.5 * height
        x2 = center_x + 0.5 * width
        y2 = center_y + 0.5 * height
        boxes_c = torch.stack([x1, y1, x2, y2], dim=1)
    return boxes_c.reshape(boxes.shape)


# 非极大值抑制
def non_max_suppress(true_boxes, confidences, iou_threshold,
                     device):  # true_boxes shape: [:,4], yolo: conf: [:,1] ; ssd: prob: [:,c] c是目标类的数量

    base = torch.argsort(confidences, descending=True)
    L = []
    while base.numel() > 1:
        L.append(base[0].item())
        iou_table = iou(true_boxes[base[0:1]], true_boxes[base[1:]])
        keep_indexes = torch.nonzero(iou_table[0, :] <= iou_threshold).reshape(-1)
        base = base[keep_indexes + 1]
    if base.numel() == 1:
        L.append(base[0].item())

    return torch.tensor(L, device=device)


def anchor_true_map(anchor_boxes, true_boxes, device, iou_threshold=0.5):  # 锚框和真实框的绑定图
    n = anchor_boxes.shape[0]
    m = true_boxes.shape[0]
    iou_table = iou(anchor_boxes, true_boxes)
    output_map = torch.full((n,), -1, dtype=torch.long, device=device)
    for _ in range(m):
        index = torch.argmax(iou_table)

        i = index.div(iou_table.shape[1], rounding_mode='trunc')
        j = index % iou_table.shape[1]
        iou_table[i, :] = -1
        iou_table[:, j] = -1
        output_map[i] = j
    for _ in range(n):
        if output_map[_] != -1:
            continue

        j = torch.argmax(iou_table[_])
        output_map[_] = j if iou_table[_][j] >= iou_threshold else -1

    return output_map


def target_bbox(anchors_boxes, labels_boxes, device):
    boxes, classes, masks = [], [], []
    n = anchors_boxes.shape[0]
    for i in range(labels_boxes.shape[0]):
        label_boxes = labels_boxes[i, :, :]
        output_map = anchor_true_map(anchors_boxes, label_boxes[:, 1:], device)
        class_labels = torch.zeros(n, dtype=torch.long, device=device)
        true_boxes = torch.zeros((n, 4), dtype=torch.float32, device=device)
        index = torch.squeeze(torch.nonzero(output_map >= 0))
        assigned_indexes = output_map[index]
        class_labels[index] = label_boxes[assigned_indexes, 0].long() + 1
        true_boxes[index] = label_boxes[assigned_indexes, 1:]
        #         class_labels = torch.max(class_labels.reshape(s*s,b),dim=1).values

        boxes.append(true_boxes)
        classes.append(class_labels)
        masks.append((output_map >= 0))

    boxes = torch.stack(boxes)
    classes = torch.stack(classes)
    masks = torch.stack(masks)

    # 真实框，每个格子包含的目标类，掩码
    return boxes, classes, masks


def mean_ap(confidences, target_true_boxes, target_predict_boxes, current_batch_size):
    mean_ap = []
    for i in range(current_batch_size):
        confidence = confidences[i]
        target_true_box = target_true_boxes[i]
        predict_box = target_predict_boxes[i]

        iou_table = iou(target_true_box, predict_box, mode='union')

        sorted_confidence_indices = confidence.sort(dim=0).indices
        ap = 0.0  # average precision
        tp_acc = 0  # 累加的tp,tp:是正类且被预测到是正类
        fp_acc = 0  # 累加的fp,fp:是负类但被预测到是正类
        p = sorted_confidence_indices.shape[0]  # tp + fn, fn:是正类，但被预测到负类, tp + fn就是正类数量之和

        for _ in range(sorted_confidence_indices.shape[0]):
            indic = sorted_confidence_indices[_]
            iou_table_specified = iou_table[indic]
            if iou_table_specified > 0.5:  # 预测框和真实框的iou大于0.5即被分为tp
                tp_acc += 1
            else:
                fp_acc += 1
            precision = tp_acc / (tp_acc + fp_acc)
            recall = tp_acc / p
            ap += recall * precision
        mean_ap.append(ap)
    return torch.tensor(mean_ap).mean()
